from llama_index import SimpleDirectoryReader, \
    GPTVectorStoreIndex, \
    PromptHelper, \
    LLMPredictor, \
    ServiceContext
from langchain.chat_models import ChatOpenAI
import gradio as gr
import os

####
# LangChain+ChatGPT 简单RAG
####

os.chdir('/')  # 文件路径
os.environ["OPENAI_API_KEY"] = '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'


def construct_index(directory_path):
    max_input_size = 4096
    num_outputs = 2000
    max_chunk_overlap = 0.5
    chunk_size_limit = 600
    prompt_helper = PromptHelper(max_input_size, num_outputs, max_chunk_overlap, chunk_size_limit=chunk_size_limit)
    llm_predictor = LLMPredictor(llm=ChatOpenAI(model_name='gpt-3.5-turbo'))
    documents = SimpleDirectoryReader(directory_path).load_data()
    service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor)
    index = GPTVectorStoreIndex.from_documents(documents, service_context=service_context)
    index.save_to_disk('index.json')
    return index


def chatbot(input_text):
    index = GPTVectorStoreIndex.load_from_disk('index.json')
    response = index.query(input_text, response_mode="compact")
    return response.response


iface = gr.Interface(fn=chatbot,
                     inputs=gr.inputs.Textbox(lines=7, label="输入您的文本"),
                     outputs="text",
                     title="AI 知识库聊天机器人")

index = construct_index("docs")
iface.launch(share=True)
